import pickle
import time
import numpy as np


from ezfranka.sim.client import RobotClient, Matrix4x4
from ezgl import Matrix4x4, Quaternion

np.set_printoptions(precision=3, suppress=True, linewidth=200)

if __name__ == '__main__':
    data = pickle.load(open(r"J:\tactile_grasp\data\episode_data.pkl", "rb"))

    robot = RobotClient()
    robot.clear_items()

    for record in data:
        ee_pose = np.array(record['ee_pose'])
        ee_pose = Matrix4x4.fromVector6d(ee_pose[0], ee_pose[1], ee_pose[2], ee_pose[3]*180/np.pi, ee_pose[4]*180/np.pi, ee_pose[5]*180/np.pi)
        
        vive_pose = record['vive_pos']
        vive_pose = Matrix4x4.fromVector6d(*vive_pose[0:3], vive_pose[5]*180/np.pi, vive_pose[4]*180/np.pi, vive_pose[3]*180/np.pi, )

        joint_angles_rad = record['joint_angles']
        
        robot.set_joint([
            joint_angles_rad[0],
            joint_angles_rad[1],
            joint_angles_rad[2],
            joint_angles_rad[3],
            joint_angles_rad[4],
            joint_angles_rad[5],
        ])
        robot.add_pose_axis(ee_pose)
        # robot.add_pose_axis(vive_pose)
        
        ee_pose_inner = robot.get_pose()
        
        # print(ee_pose.toVector6d()[3:] * np.pi / 180)
        # print(ee_pose_inner.toVector6d()[3:] * np.pi / 180)
        print(ee_pose_inner.toVector6d()[3:] * np.pi / 180 - ee_pose.toVector6d()[3:] * np.pi / 180)
        # break
        time.sleep(2)

        